import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from utils.utils import _gather_feature, _tranpose_and_gather_feature, flip_tensor,draw_hist




def _softmax(x):
    y = nn.Softmax(dim=2)(x)
    return y

def _nms(heat, kernel=3):
    hmax = F.max_pool2d(heat, kernel, stride=1, padding=(kernel - 1) // 2)
    keep = (hmax == heat).float()
    return heat * keep
    
    
    
def _topk(scores, K=40):
    batch, cat, height, width = scores.size()

    topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)
#     print("#####################",topk_scores)
    
    topk_inds = topk_inds % (height * width)
    
#     topk_ys = (topk_inds / width).int().float()#torch 1.5版本不再有这种除法
    topk_ys = torch.floor_divide(topk_inds,width).int().float()
#     result = torch.floor_divide(A, n)
    
    topk_xs = (topk_inds % width).int().float()

    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
#     print("---------------------------",topk_score)
    
#     topk_clses = (topk_ind / K).int()
    topk_clses = torch.floor_divide(topk_ind,K).int()
    
    topk_inds = _gather_feature(topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_ys = _gather_feature(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_xs = _gather_feature(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_clses, topk_ys, topk_xs


def ctdet_decode(hmap, regs, w_h_, K=100):
    batch, cat, height, width = hmap.shape
    hmap = torch.sigmoid(hmap)

    # if flip test
    if batch > 1:
        hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
        w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2
        regs = regs[0:1]

    batch = 1

    hmap = _nms(hmap)  # perform nms on heatmaps

#     hmap = _nms_me(hmap)

    scores, inds, clses, ys, xs = _topk(hmap, K=K)

    regs = _tranpose_and_gather_feature(regs, inds)
    regs = regs.view(batch, K, 2)
#     print("regs",regs)
    xs = xs.view(batch, K, 1) + regs[:, :, 0:1]
    ys = ys.view(batch, K, 1) + regs[:, :, 1:2]

    w_h_ = _tranpose_and_gather_feature(w_h_, inds)
    w_h_ = w_h_.view(batch, K, 2)

    clses = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)

    bboxes = torch.cat([xs - w_h_[..., 0:1] / 2,
                        ys - w_h_[..., 1:2] / 2,
                        xs + w_h_[..., 0:1] / 2,
                        ys + w_h_[..., 1:2] / 2], dim=2)
    detections = torch.cat([bboxes, scores, clses], dim=2)
    return detections




def ctdet_decode_tblr(hmap, tblr, K=100):
    batch, cat, height, width = hmap.shape
    hmap = torch.sigmoid(hmap)

    # if flip test
#     if batch > 1:
#         hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
#         w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2


    batch = 1

    hmap = _nms(hmap)  # perform nms on heatmaps
#     print(hmap)
    scores, inds, clses, ys, xs = _topk(hmap, K=K)
#     print("sssssssssssssssssssssssss",scores)
    tblr = _tranpose_and_gather_feature(tblr, inds)
#     regs = regs.view(batch, K, 2)
#     print("regs",regs)
    xs = xs.view(batch, K, 1) 
    ys = ys.view(batch, K, 1)
    tblr = tblr.view(batch, K, 4)
#     print(tblr.shape)
    clses = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)
#     print("gggggggggggggggggggggggggggggg",scores)

    #神坑！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！(...)三个点
    bboxes = torch.cat([xs - tblr[...,2:3],
                        ys - tblr[...,0:1],
                        xs + tblr[...,3:4],
                        ys + tblr[...,1:2]], dim=2)
#     print(bboxes)
    detections = torch.cat([bboxes, scores, clses], dim=2)
#     print("000000000000000000000000000000000000000",detections)
    return detections



def ctdet_decode_dfl(hmap, tblr,reg_max, K=100):
    project = torch.linspace(0, reg_max, reg_max + 1)
    
    batch, cat, height, width = hmap.shape
    hmap = torch.sigmoid(hmap)

    # if flip test
#     if batch > 1:
#         hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
#         w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2


    batch = 1

    hmap = _nms(hmap)  # perform nms on heatmaps
#     print(hmap)
    scores, inds, clses, ys, xs = _topk(hmap, K=K)
#     print("sssssssssssssssssssssssss",scores)
    tblr = _tranpose_and_gather_feature(tblr, inds)
#     print(tblr.shape)
    t = tblr[:,:,:reg_max+1]
    b = tblr[:,:,reg_max+1:2*reg_max+2]
    l = tblr[:,:,2*reg_max+2:3*reg_max+3]
    r = tblr[:,:,3*reg_max+3:4*reg_max+4]
    t_ = Integral(t,project)
    b_ = Integral(b,project)
    l_ = Integral(l,project)
    r_ = Integral(r,project)
#     if np.random.random() < 0.1:
#         print(t_,b_,l_,r_)
    
#     print(t_.shape)
#     print(t_,b_,l_,r_)
#     tblr = 
#     tblr = torch.tensor([t_,b_,l_,r_])
#     regs = regs.view(batch, K, 2)
#     print("regs",regs)
#     xs = xs.view(batch, K, 1) 
#     ys = ys.view(batch, K, 1)
#     tblr = tblr.view(batch, K, 4)
#     print(tblr.shape)
    clses = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)
#     print("gggggggggggggggggggggggggggggg",scores)

    #神坑！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！(...)三个点
    bboxes = torch.cat([xs - l_,
                        ys - t_,
                        xs + r_,
                        ys + b_], dim=2)
#     print(bboxes)
    detections = torch.cat([bboxes, scores, clses], dim=2)
#     print("000000000000000000000000000000000000000",detections)
    return detections

def Integral(x,project):
    weight = x
#     weight = F.softmax(x, dim=2)
    pred_ = weight*project.type_as(weight)
    pred_final = torch.unsqueeze(torch.sum(pred_,dim=2),dim=2)
    
    return pred_final
  
    



def ctdet_decode_csl_one(hmap, tblr, reg_max,omega,K=100):
    cls_num = int(reg_max/omega)+1
    project = torch.linspace(0, cls_num-1, cls_num )
    batch, cat, height, width = hmap.shape
    hmap = torch.sigmoid(hmap)

    # if flip test
#     if batch > 1:
#         hmap = (hmap[0:1] + flip_tensor(hmap[1:2])) / 2
#         w_h_ = (w_h_[0:1] + flip_tensor(w_h_[1:2])) / 2


    batch = 1

    hmap = _nms(hmap)  # perform nms on heatmaps
#     print(hmap)
    scores, inds, clses, ys, xs = _topk(hmap, K=K)
    
#     print("sssssssssssssssssssssssss",scores)
    tblr = _tranpose_and_gather_feature(tblr, inds)
#     print(tblr.shape)
    t = tblr[:,:,:cls_num]
    b = tblr[:,:,cls_num:2*cls_num]
    l = tblr[:,:,2*cls_num:3*cls_num]
    r = tblr[:,:,3*cls_num:4*cls_num]
#     regs = regs.view(batch, K, 2)
#     print("regs",regs)
    xs = xs.view(batch, K, 1) 
    ys = ys.view(batch, K, 1)
#------------------------------------------------------------------------------------------    
    t = _softmax(t.view(batch, K, cls_num))
    
    ##############################################################################
#     画图用

 ##############################################################################
#     print(t)
#     print(torch.max(t,dim=2))
    t_s =torch.max(t,dim=2)[0].view(batch, K, 1)
#     print(t)
    t = t.float()
    
    b = _softmax(b.view(batch, K, cls_num))

    b_s =torch.max(b,dim=2)[0].view(batch, K, 1)
    b = b.float()
    
    l = _softmax(l.view(batch, K, cls_num))

    l_s =torch.max(l,dim=2)[0].view(batch, K, 1)
    l = l.float()

    r = _softmax(r.view(batch, K, cls_num))

    r_s =torch.max(r,dim=2)[0].view(batch, K, 1)
    r = r.float()
#     print(K)
#     print(t.shape)
#------------------------------------------------------------------------------------------ 
#     print(tblr.shape)

#     reg_score = torch.floor_divide(t_s+b_s+l_s+r_s,4)
###############################################################################
    #画图用
#     t_list = t[0,0,:].cpu().data.numpy()
#     print(t[0,0,:].cpu().data.numpy())
#     draw_hist(t_list, 'Top', 'Distance', 'Value', 0, reg_max, 0.0, 1.0,"t")  # 直方图展示
    
#     b_list = b[0,0,:].cpu().data.numpy()
#     print(b[0,0,:].cpu().data.numpy())
#     draw_hist(b_list, 'Bottom', 'Distance', 'Value', 0, reg_max, 0.0, 1.0,"b")  # 直方图展示
    
#     l_list = l[0,0,:].cpu().data.numpy()
#     print(l[0,0,:].cpu().data.numpy())
#     draw_hist(l_list, 'Left', 'Distance', 'Value', 0, reg_max, 0.0, 1.0,"l")  # 直方图展示
    
#     r_list = r[0,0,:].cpu().data.numpy()
#     print(r[0,0,:].cpu().data.numpy())
#     draw_hist(r_list, 'Right', 'Distance', 'Value', 0, reg_max, 0.0, 1.0,"r")  # 直方图展示
 ###############################################################################
    t = Integral(t,project)*omega
    b = Integral(b,project)*omega
    l = Integral(l,project)*omega
    r = Integral(r,project)*omega
#     t = t*2
#     b = b*2
#     l = l*2
#     r = r*2
    clses = clses.view(batch, K, 1).float()
    scores = scores.view(batch, K, 1)
    
#     scores = scores*reg_score
#     print("gggggggggggggggggggggggggggggg",scores)

    #神坑！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！！(...)三个点
    bboxes = torch.cat([xs - l,
                        ys - t,
                        xs + r,
                        ys + b], dim=2)
#     print(bboxes)
    detections = torch.cat([bboxes, scores, clses], dim=2)
#     print(detections.shape)
#     print("000000000000000000000000000000000000000",detections)
    return detections

